import psutil
import time
import threading
from typing import Optional
import argparse

class CPUController:
    def __init__(self):
        self.running = True
        self.target_usage = 50
        self.control_interval = 0.1  # 控制间隔，单位秒
        self.cpu_count = psutil.cpu_count()
        self.target_core: Optional[int] = None
    
    def busy_wait(self, duration):
        """执行忙等待"""
        start_time = time.time()
        while time.time() - start_time < duration:
            pass
    
    def control_loop(self):
        """CPU使用率控制循环"""
        while self.running:
            # 获取当前CPU使用率
            if self.target_core is not None:
                current_usage = psutil.cpu_percent(interval=0.1, percpu=True)[self.target_core]
            else:
                current_usage = psutil.cpu_percent(interval=0.1)
            
            # 根据当前使用率调整工作时间
            if current_usage < self.target_usage:
                self.busy_wait(self.control_interval * 0.8)
            else:
                time.sleep(self.control_interval * 0.2)
    
    def start(self, target_usage: float, target_core: Optional[int] = None):
        """启动CPU控制"""
        self.target_usage = target_usage
        self.target_core = target_core
        
        if target_core is not None and (target_core < 0 or target_core >= self.cpu_count):
            raise ValueError(f"Invalid core number. Must be between 0 and {self.cpu_count-1}")
        
        thread = threading.Thread(target=self.control_loop)
        thread.daemon = True
        thread.start()
        return thread

def main():
    parser = argparse.ArgumentParser(description='Control CPU usage')
    parser.add_argument('--usage', type=float, default=50,
                      help='Target CPU usage percentage (default: 50)')
    parser.add_argument('--core', type=int, default=None,
                      help='Target CPU core number (optional)')
    
    args = parser.parse_args()
    
    controller = CPUController()
    try:
        thread = controller.start(args.usage, args.core)
        print(f"控制已启动: 目标使用率 {args.usage}%"
              f"{f' 在核心 {args.core}' if args.core is not None else ''}")
        print("按 Ctrl+C 停止...")
        thread.join()
    except KeyboardInterrupt:
        controller.running = False
        print("\n程序已停止")

if __name__ == "__main__":
    main() 